查看原文
其他

华为GTS LocMoE+:高可扩展性亲和度 MoE 架构,低开销实现主动路由

机器之心
2024-09-06

AIxiv专栏是机器之心发布学术、技术内容的栏目。过去数年,机器之心AIxiv专栏接收报道了2000多篇内容,覆盖全球各大高校与企业的顶级实验室,有效促进了学术交流与传播。如果您有优秀的工作想要分享,欢迎投稿或者联系报道。投稿邮箱:liyazhou@jiqizhixin.com;zhaoyunfeng@jiqizhixin.com

本论文共同一作为李婧博士,孙志杰和林大超博士,主要成员来自GTS AI计算Lab,主要研究及落地领域包含LLM训推加速、AI训练保障和图计算。


MoE 因其在训推流程中低销高效的特点,近两年在大语言模型领域大放异彩。作为 MoE 的灵魂,专家如何能够发挥出最大的学习潜能,相关的研究与讨论层出不穷。此前,华为 GTS AI 计算 Lab 的研究团队提出了 LocMoE ,包括新颖的路由网络结构、辅助降低通信开销的本地性 loss 等,引发了广泛关注。

LocMoE 的上述设计有效缓解了部分 MoE 经典结构在训练中的瓶颈,如: 专家路由算法可能无法有效区分 token,以及通信同步效率受限于节点内和节点间的传输带宽差异等。此外,LocMoE 证明并求解出在能够成功处理判别性 token 的前提下专家容量的下限。该下限是根据 token 被动派发给专家的场景下,判别性 token 存在于 token 批次中的概率分布而得出。那么,如果专家同样具备择优 token 的能力,判别性 token 被处理的概率将大幅提高,专家容量下限将能进一步压缩。

在上述设想的基础上,该团队进一步提出了基于低开销主动路由的 MoE 架构,将其命名为 LocMoE+。LocMoE+ 继承了 LocMoE 高判别性专家以及本地性通信优势,进一步改造路由策略,定义 token 与专家之间的亲和度指标,并以此指标入手,更有效率完成 token 分派,进而提升训练效率。


  • 论文链接:https://arxiv.org/pdf/2406.00023

论文简介

该论文的核心思路是通过结合传统被动路由与专家主动路由,提升一定容量下处理判别性 token 的概率,从而降低样本噪声和提升训练效率。该论文从 token 和其被分派的专家之间的关系入手,以一种低计算开销的方案量化和定义了专家与 token 之间的亲和性。据此,该论文实现了全局的自适应路由策略,并基于亲和性得分以专家维度对 token 进行重排和选择。同时,专家容量的下限被证明随着 token 特征分布的稳定而逐渐减小,训练开销得以降低。

该论文率先结合两种路由机制,根据学习性路由策略中 token 倾向于路由至与其夹角更小的专家的发现,打破了主动路由现存方案开销过大影响训练效率的障碍,并与被动路由的本质保持统一。

值得一提的是,作者选用了和 LocMoE 工作完全不同的硬件环境(服务器型号,NPU卡型号,集群组网方案)、训练框架和骨干模型,以证明该系列工作的高扩展性和易于移植性。

自适应双向路由分派机制

背景介绍

传统 MoE 有两种路由分派机制:

(1)hard router,直接把整个 token 特征进行分派;
(2)soft router,把 token 特征的加权组合进行分派。

本文延续考虑(1)因其较低的计算代价。而对于 hard router 场景,又可分为 1)Token Choice Router (TCR), 也就是让每个 token 去选 top-k 专家;2)Expert Choice Router (ECR),也就是让每个专家去选择 top-C 合适的 token。由于容量限制,每个专家接收的 token 数有上限 C, 因此在场景 1)中,会对每个专家接收的 token进行截断:


前人工作指出 MoE 训练分为两阶段:Phase 1. 路由训练,保障路由可以合理分派 token,也就是能把不同领域或者差异大的 token 进行区分和分派给不同专家;Phase 2. 由于 token 路由的作用,每个专家接收到相同领域或者性质相近的 token,每个专家在经历一定训练后就可以习得相关领域和性质的知识。总结来说,MoE 每步训练“成功”的关键在于 token 分派的正确和合理性。

该文贡献

(1)通过 softmax 激活函数推演,专家和 token 之间的余弦相似度能够较为准确地衡量亲和度:

 
(2)从理论建模角度,分析 TCR 和 ECR 两个常见场景下单次训练成功率:
 


基于理论,作者指出

  • 在模型训练初期,路由分派 token 能力不足时,每次训练 TCR 要比 ECR 以更高概率成功训练,而且需要较大的专家容量保证选到合适的 token。
  • 在模型训练后期,路由有一定能力正确分派 token 时,每次训练 ECR 要比 TCR 以更高概率成功训练,此时只需要较小的容量就可以选到合适的 token。
 


该理论也很符合直观,路由没有分派能力时,不如让 token 随机选择专家,而当路由有一定分派能力,也就是专家能选到合适 token 的时候,使用 ECR 更合适。因此作者推荐 TCR 向 ECR 方式的转变,提出了全局级别的自适应路由切换策略,同时基于专家容量的需求估计,在训练后期使用更小的专家容量。

实验结果

该论文的实验在昇腾 910B3 NPU 的自建集群上进行,得益于华为专有的高速缓存一致性协议高性能计算系统(HCCS)在多卡场景下实现的高性能设备间数据通信,以及专为昇腾处理器设计的华为集合通信库(HCCL)在 HCCS 等高速链路上实现高性能分布式训练。实验采用兼容昇腾 NPU 的 PyTorch for Ascend 框架和专为昇腾设备定制的加速库 AscendSpeed 和训练框架 ModelLink,关注 LLM 并行策略和通信遮掩优化。

训练效率

实验结果表明,在不影响模型训练收敛性或有效性的情况下,每个专家至少需要处理的token数量较基线可以减少60% 以上。结合通信优化,在 32 卡、64 卡和 256 卡的集群规模下,训练效率平均提高 5.4% 至 46.6%。

 
显存占用

LocMoE+ 对于显存占用也存在一定增益,尤其是集群规模较小计算较密集的场景。使用 Ascend Insight 工具对显存监控样本进行分析,可以看出,LocMoE+ 显存使用相比于基线下降了 4.57% 至 16.27%,相比于 LocMoE 下降了 2.86% 至 10.5%。随着集群规模的增大,显存占用的差距随之缩小。

 
效果评测

开源评测集 C-Eval 和 TeleQnA,以及自主构建的 ICT 领域评测集 GDAD 被用于评估 LocMoE+ 在通用知识和领域知识上的能力。其中,GDAD 共计覆盖了 47 个子项,包含 18060 条样本,考察模型在领域任务,领域能力认证考试和通用能力三大评估体系中的表现。

经过充分 SFT 后,LocMoE+ 相比于基线在领域任务能力的 16 个子能力上平均提升了约 20.1%,相比于 LocMoE 则提升了约3.5%。领域能力认证考试则分别提升了 16% 和 4.8%。在通用能力的 18 个子能力中,LocMoE+ 分别提升约 13.9% 和 4.8%。整体而言,LocMoE+ 在 GDAD、C-Eval 和 TeleQnA 上分别表现出 9.7% 到 14.1% 的效果提升。


 


© THE END 

转载请联系本公众号获得授权

投稿或寻求报道:content@jiqizhixin.com

继续滑动看下一个
机器之心
向上滑动看下一个

您可能也对以下帖子感兴趣

文章有问题?点此查看未经处理的缓存